
#
# AES Counter Mode (using pycrypto) to interoperate with the aes-lib.js implementation by Chris Veness
# http://www.movable-type.co.uk/scripts/aes.html
#
# Copyright (C) 2009 Grahame Bowland <grahame@angrygoats.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
# 
# You should have received a copy of the GNU Lesser General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#


from Crypto.Cipher import AES
import base64, math, itertools, datetime, struct

class CTR(object):
    blocksize = 16
    
    @classmethod
    def urandom_8b(self):
        fd = open('/dev/urandom')
        while True:
            bytes = fd.read(8)
            yield bytes
        fd.close()
    
    @classmethod
    def ord_str(cls, s):
        return [ord(t) for t in s]

    @classmethod
    def chr_ord(cls, s):
        return ''.join([chr(t) for t in s])
            
    @classmethod
    def pad(cls, string, try_lens):
        for l in try_lens:
            if len(string) <= l:
                return string + ('\0' * (l - len(string)))
        return string[:l]
    
    @classmethod
    def blocks(cls, s):
        nblocks = int(math.ceil(len(s) / (1.*CTR.blocksize)))
        for i in xrange(nblocks):
            yield s[i*CTR.blocksize:(i+1)*CTR.blocksize]
    
    @classmethod
    def generate_counterblock(cls, block_num, nonce):
        counterblock = nonce + []
        for i in xrange(7,-1,-1):
            v = (block_num / (2 ** (8*i))) & 0xFF
            counterblock.append(v)
        return counterblock
        
    def __init__(self, password, key_length):
        if key_length != 128 and key_length != 192 and key_length != 256:
            raise Exception("invalid key length, must be 128, 192 or 256")
        self.nbytes = key_length / 8
        pwbytes = CTR.pad(password, (self.nbytes,))
        cipher = AES.new(pwbytes)
        self.key = cipher.encrypt(pwbytes)
        self.urandom = CTR.urandom_8b()

    def ctr(self, nonce, text):
        res = []
        cipher = AES.new(self.key)
        for block_num, block in enumerate(CTR.blocks(text)):
            counterblock = CTR.generate_counterblock(block_num, nonce)
            ciphered = cipher.encrypt(''.join((chr(t) for t in counterblock)))
            res += (x ^ y for (x, y) in itertools.izip(CTR.ord_str(ciphered), CTR.ord_str(block)))
        return res

    def decrypt(self, text, return_nonce=False):
        text = base64.decodestring(text)
        nonce = CTR.ord_str(text[0:8])
        dec = CTR.chr_ord(self.ctr(nonce, text[8:]))
        if return_nonce:
            return dec, nonce
        else:
            return dec

    def generate_nonce(self):
        u = self.urandom.next()
        n = struct.pack('hbbbbbbh', *datetime.datetime.now().timetuple()[:-1])
        return [ ord(x) ^ ord(y) for (x, y) in itertools.izip(u, n) ]
    
    def encrypt(self, text, nonce=None):
        cipher = AES.new(self.key)
        if nonce is None:
            nonce = self.generate_nonce()
        res = nonce + self.ctr(nonce, text)
        return base64.encodestring(CTR.chr_ord(res)).replace('\n', '')
